# 全角度
import mujoco
import control as ct
import numpy as np
import mujoco.viewer as viewer
from os.path import dirname, join

# 加载模型
model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'cartpole.xml')) 
data = mujoco.MjData(model)


M = model.body_mass[1]  # 小车的质量
m = model.body_mass[2]  # 摆杆的质量
g = - model.opt.gravity[2]  # 重力加速度
L = model.geom_size[2][1] # 摆杆的长度

'''
y_1 = x
y_2 = \dot{x}
y_3 = \theta
y_4 = \dot{\theta}
'''
A = np.array([[0, 1, 0, 0,], [0, 0, -3*m/(4*M+m)*g, 0],[0, 0, 0, 1],[0, 0, 6*(M+m)/((4*M+m)*L)*g, 0]])
B = np.array([[0],[4/(4*M+m)], [0], [-6/((4*M+m)*L)]])
Q = np.identity(4)
R = np.identity(1)
print(f"A = {A}\n\n B = \n{B}\n\n Q = \n{Q}\n\n R = \n{R}\n\n")
K, S, E = ct.lqr(A, B, Q, R)
print(f"K = \n{K}\n\n S = \n{S}\n\n E = \n{E}\n\n")

def normalize_angles(theta):
    return np.remainder(theta + np.pi, 2 * np.pi) - np.pi

def compute_F(y1, y2, y3, y4, M, m, L, g, delta):
    # 计算公共分母项 (6L y4 - 4y2)
    denominator = 6 * L * y4 - 4 * y2
    
    # 计算分子部分
    term1 = 3 * m * (L**2 * y4**3 - g * y2) * np.sin(y3) * np.cos(y3)
    term2 = 2 * L * y4 * (3 * M * g + 3 * m * g - m * y2 * y4) * np.sin(y3)
    term3 = (y1 * y2 + L**2 * y3 * y4) * (4 * M + m + 3 * m * np.sin(y3)**2)
    numerator = term1 + term2 + term3
    
    # 分段计算 F
    if abs(denominator) > delta:
        # 情况1：分母绝对值较大时直接计算
        F = numerator / denominator + denominator
    else:
        # 情况2：分母接近零时使用符号函数平滑处理
        sign_denominator = np.sign(denominator)
        F = numerator / (sign_denominator * delta) + denominator
    return F

def myControl(model, data):
    x_raw = data.sensordata
    x = x_raw
    x[2] = normalize_angles(x_raw[2])
    xref = np.array([0, 0, 0, 0])
    x_error = x-xref
    if x[2] < np.pi/6 and x[3] < 1:
        u = -K[0]@x_error
    else:
        u = compute_F(x[0], x[1], x[2], x[3], M, m, L, g, delta = 0.1)
    data.ctrl = u
    return u

mujoco.mj_resetDataKeyframe(model, data, 0)
# 运行仿真
mujoco.set_mjcb_control(myControl)
viewer.launch(model,data)